package com.wsjj.yjh;

import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.streaming.api.datastream.ConnectedStreams;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.CoProcessFunction;
import org.apache.flink.util.Collector;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class ConnectKeybyDemo {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        DataStreamSource<Tuple2<Integer, String>> source1 = env.fromElements(
                Tuple2.of(1, "a1"),
                Tuple2.of(1, "a2"),
                Tuple2.of(2, "b"),
                Tuple2.of(3, "c")
        );
        DataStreamSource<Tuple3<Integer, String, Integer>> source2 = env.fromElements(
                Tuple3.of(1, "aa1", 1),
                Tuple3.of(1, "aa2", 2),
                Tuple3.of(2, "bb", 1),
                Tuple3.of(3, "cc", 1)
        );
        ConnectedStreams<Tuple2<Integer, String>, Tuple3<Integer, String, Integer>> connect = source1.connect(source2);

        ConnectedStreams<Tuple2<Integer, String>, Tuple3<Integer, String, Integer>> tuple2Tuple3ConnectedStreams = connect.keyBy(new KeySelector<Tuple2<Integer, String>, Integer>() {
            @Override
            public Integer getKey(Tuple2<Integer, String> value) throws Exception {
                return value.f0;
            }
        }, new KeySelector<Tuple3<Integer, String, Integer>, Integer>() {
            @Override
            public Integer getKey(Tuple3<Integer, String, Integer> value) throws Exception {
                return value.f0;
            }
        });
        tuple2Tuple3ConnectedStreams.process(new CoProcessFunction<Tuple2<Integer, String>, Tuple3<Integer, String, Integer>, String>() {

           Map<Integer, List<Tuple2<Integer,String>>> cache1  = new HashMap<>();
           Map<Integer,List<Tuple3<Integer,String,Integer>>> cache2  = new HashMap<>();
            @Override
            public void processElement1(Tuple2<Integer, String> value, CoProcessFunction<Tuple2<Integer, String>, Tuple3<Integer, String, Integer>, String>.Context ctx, Collector<String> out) throws Exception {
                Integer f0 = value.f0;
                if(!cache1.containsKey(f0)){
                    List<Tuple2<Integer, String>> tuple1s = new ArrayList<>();
                    tuple1s.add(value);
                    cache1.put(f0,tuple1s);
                }else {
                    cache1.get(f0).add(value);
                }
                if(cache2.containsKey(f0)){
                    for (Tuple3<Integer, String, Integer> integerStringIntegerTuple3 : cache2.get(f0)) {
                        out.collect("s1:" + value + "<--------->s2:" + cache2);
                    }
                }
            }

            @Override
            public void processElement2(Tuple3<Integer, String, Integer> value, CoProcessFunction<Tuple2<Integer, String>, Tuple3<Integer, String, Integer>, String>.Context ctx, Collector<String> out) throws Exception {

                Integer f0 = value.f0;
                if(!cache2.containsKey(f0)){
                    ArrayList<Tuple3<Integer, String, Integer>> tuple3s = new ArrayList<>();
                    tuple3s.add(value);
                    cache2.put(f0,tuple3s);
                }else {
                    cache2.get(f0).add(value);
                }
                if(cache1.containsKey(f0)){
                    for (Tuple2<Integer, String> integerStringTuple2 : cache1.get(f0)) {
                        out.collect("s2:" + value + "<--------->s1:" + cache1);
                    }
                }
            }
        }).print();
        env.execute();
    }
}
